
#include <string>
#include <fstream>
#include <iomanip>
#include <algorithm>
#include <math.h>
#include <string>

#include "DcqcnRoceConnection.h"
#include "RoceMsg_m.h"
#include "cluster/messages/messages_m.h"
#include "inet/networklayer/common/EcnTag_m.h"

namespace ai_sim {

Define_Module(DcqcnRoceConnection);

void DcqcnRoceConnection::init(const inet::L3Address& src_addr, 
        const inet::L3Address& dst_addr,
        int src_port, 
        int dst_port) 
{
    qp.srcAddr = src_addr;
    qp.dstAddr = dst_addr;
    qp.srcPort = src_port;
    qp.dstPort = dst_port;
    qp.type = RoceType::DCQCN;
}

uint32_t DcqcnRoceConnection::getTransDelayNs(uint32_t sizeInBytes, int srdIdx, int destIdx) {
    //
    ASSERT(sizeInBytes != 0);
    // rate is Mbps
    if((int)dcqcn.curRate != preRate) {
        // if(getParentModule()->getParentModule()->getParentModule()->getParentModule()->getIndex() == 0){
        //     std::cout<<srdIdx<<"-"<<destIdx<<"-"<< dcqcn.curRate<<"-"<<omnetpp::simTime()<<endl;
        // }
        char* pc =new char[100];//足够长
        strcpy(pc,getParentModule()->getParentModule()->getParentModule()->getParentModule()->getName());
        std::ofstream ofile;
        ofile.open("./record_rate.txt",ios::app);
        ofile << srdIdx << "-" << destIdx << "-" << dcqcn.curRate << "-" << omnetpp::simTime() <<endl;
        ofile.close();
        preRate = dcqcn.curRate;
    }
    return (uint32_t)( sizeInBytes * 8 / dcqcn.curRate * 1e3 );
}


/****************************
 *  /// 
 ****************************/
// GPU send message to Connection
void DcqcnRoceConnection::HandlePacketFromHigherLayer(const FlowMsg& f_msg) {
    senderMgr.putLargeChunk(f_msg.getThe_flow());
}

// packet from socket, may be ACK, NACK, CNP, DATA.
CLUSTER_MSGS DcqcnRoceConnection::HandlePacketFromLowerLayer(inet::Packet& packet) {
    //
    auto roceMsg =  packet.peekAtFront<RoceMsg>();
    uint32_t seqNum = roceMsg->getSequenceNum();
    uint32_t msgKind = roceMsg->getKind();
    const Flow& flow = roceMsg->getThe_flow();
    EV << "In roce connection\nFlow id: " << flow.flowId << "\n";
    bool lastSeg = roceMsg->isLastSegment();
    if (msgKind == CLUSTER_MSGS::SOCK_TO_SOCK_DATA_MSG) {
        Flow reply_flow;
        reply_flow.flowId = flow.flowId;
        reply_flow.data_size = 1;
        reply_flow.srcIdx = flow.destIdx;
        reply_flow.destIdx = flow.srcIdx;
        reply_flow.outportId = flow.outportId;
        int ecnind = packet.findTag<inet::EcnInd>()->getExplicitCongestionNotification();
        if (ecnind == 3) { // may generate cnp
            //
            if (simTime() > cnpSendInfo.lastCnpSentTime + cnpSendInfo.cnpGap) {
                auto cnpPacket = wrapPacket(1, CLUSTER_MSGS::CNP_MSG, -1, false, reply_flow);
                ctrlQueue.push(cnpPacket);
                cnpSendInfo.lastCnpSentTime = simTime();
            } else {
                EV << "dont generate CNP";
            }
        }
        int ackOrNack = receiverMgr.processReceivedPacket(seqNum, lastSeg, flow);
        if (ackOrNack == 1) {
            auto ackPacket = wrapPacket(1, CLUSTER_MSGS::ACK_MSG, receiverMgr.getAckSeqNum(), false, reply_flow);
            ctrlQueue.push(ackPacket);
        } else if (ackOrNack == 2) {
            // generate NACK
            // out-of-order delivery
            // perform GBN logic
            auto nackPacket = wrapPacket(1,  CLUSTER_MSGS::NACK_MSG, receiverMgr.getAckSeqNum(), false, reply_flow);
            ctrlQueue.push(nackPacket);
        }
        return CLUSTER_MSGS::SOCK_TO_SOCK_DATA_MSG;
    } else if (msgKind == CLUSTER_MSGS::ACK_MSG) {
        handleAck(seqNum);
        return CLUSTER_MSGS::ACK_MSG;
    } else if (msgKind == CLUSTER_MSGS::NACK_MSG) {
        handleNack(seqNum);
        return CLUSTER_MSGS::NACK_MSG;
    } else if (msgKind == CLUSTER_MSGS::CNP_MSG) {
        handleCnp();
        return CLUSTER_MSGS::CNP_MSG;
    } else {
        error("Unrecognized Roce message type.");
    }
    return CLUSTER_MSGS::UNKONWN_MSG;
}

inet::Packet* DcqcnRoceConnection::GetNextDataPacket(simtime_t* time) {
    if (dataSendInfo.lastSentTime + dataSendInfo.dataGap > simTime()) {
        // please wait, now allow send
        return nullptr;
    }
    if (senderMgr.isSendCompleted()) {
        return nullptr;
    }
    // Make data packet
    SegmentInfo* segInfo = senderMgr.getNextSegment();
    if(segInfo == nullptr) {
        return nullptr;
    }
    int sendSize = segInfo->segmentSize;
    ASSERT(sendSize != 0);
    auto packet = wrapPacket(sendSize, 
            CLUSTER_MSGS::SOCK_TO_SOCK_DATA_MSG, 
            segInfo->seqNum, 
            segInfo->isLastSegment, 
            segInfo->flow);
    // std::cout << "generate data packet\n";
    uint32_t transTime = getTransDelayNs(sendSize, segInfo->flow.srcIdx, segInfo->flow.destIdx);
    *time = SimTime(transTime, SIMTIME_NS);
    dataSendInfo.dataGap = omnetpp::SimTime(transTime, SIMTIME_NS);
    dataSendInfo.lastSentTime = simTime();
    delete segInfo;
    return packet;
}

inet::Packet* DcqcnRoceConnection::GetNextControlPacket() {
    if (ctrlQueue.empty()) {
        return nullptr;
    }
    auto packet = ctrlQueue.front();
    ctrlQueue.pop();
    return packet;
}

FlowMsg* DcqcnRoceConnection::GetReceivedFlowMessage() {
    if (receiverMgr.hasFinishedFlow()) {
        FlowMsg* flowMsg = new FlowMsg;
        flowMsg->setKind(CLUSTER_MSGS::DATA_MSG);
        flowMsg->setThe_flow(receiverMgr.getFinishedFlow());
        return flowMsg;
    }
    return nullptr;
}

/*********
 * DCQCN
 ********/

void DcqcnRoceConnection::initialize() {
    // init DCQCN
    dcqcn.curRate = LINE_RATE;
    dcqcn.additiveRate = AI_RATE;
    dcqcn.hyperIncRate = HAI_RATE;
    dcqcn.targetRate = 0;
    dcqcn.maxRate = LINE_RATE;
    dcqcn.minRate = QP_MINIMUM_RATE;
    dcqcn.alpha = 1;
    dcqcn.decreaseCnpArrived = false;
    dcqcn.alphaCnpArrived = false;
    dcqcn.firstCnpArrived = true;
    dcqcn.rateIncStage = 0;
    dcqcn.g = (double)1 / 256;
    timers.alphaTimer = nullptr;
    timers.rateIncTimer = nullptr;
    timers.rateDecTimer = nullptr;
    
    // as recommended in DCQCN paper
    // see [https://conferences.sigcomm.org/sigcomm/2015/pdf/papers/p523.pdf]
    dcqcn.alphaUpdateInter = SimTime(55, SIMTIME_US); // by default 55us
    // Debug: Modify rateIncInter, rateDecInter
    dcqcn.rateIncInter = SimTime(1500, SIMTIME_US); // by default 1.5ms
    dcqcn.rateDecInter = SimTime(4, SIMTIME_US); // by default 4us

    dataSendInfo.dataGap = SimTime(0); // for get next packet
    dataSendInfo.lastSentTime = simTime();
    // Note generate CNP is expensive operation
    // we set dataGap to 50us
    cnpSendInfo.cnpGap = SimTime(50, SIMTIME_US);
    cnpSendInfo.lastCnpSentTime = simTime();

    // Start timers
    timers.alphaTimer = new cMessage("alphaUpdate");
    timers.alphaTimer->setContextPointer(this);
    scheduleTimeout(timers.alphaTimer, dcqcn.alphaUpdateInter);

    timers.rateIncTimer = new cMessage("rateInc");
    timers.rateIncTimer->setContextPointer(this);
    scheduleTimeout(timers.rateIncTimer, dcqcn.rateIncInter);

    timers.rateDecTimer = new cMessage("rateDec");
    timers.rateDecTimer->setContextPointer(this);
    scheduleTimeout(timers.rateDecTimer, dcqcn.rateDecInter);

    dcqcn.targetRate = dcqcn.curRate;
}

void DcqcnRoceConnection::handleCnp() {
    dcqcn.alphaCnpArrived = true;
    dcqcn.decreaseCnpArrived = true;
    // std::cout << "cnp packet arrived" << endl;
    // (TODO): Uncomment the following logic after rearranging
    // the logic of RoceAdapter
    // if (dcqcn.firstCnpArrived) {
    //     std::cout << "rate decrease" << endl;
    //     // alpha update timer
    //     if (timers.alphaTimer == nullptr) {
    //         timers.alphaTimer = new cMessage("alphaUpdate");
    //         timers.alphaTimer->setContextPointer(this);
    //         scheduleTimeout(timers.alphaTimer, dcqcn.alphaUpdateInter);
    //     }
    //     // rate increase timer
    //     if (timers.rateIncTimer == nullptr) {
    //         timers.rateIncTimer = new cMessage("rateInc");
    //         timers.rateIncTimer->setContextPointer(this);
    //         // scheduleTimeout(rateIncTimer, alphaUpdateInter);
    //     }
    //     rateDecrease();
    //     dcqcn.firstCnpArrived = false;
    //     dcqcn.targetRate = dcqcn.curRate;
    // }
}

void DcqcnRoceConnection::updateAlpha() {
    //
    if (dcqcn.alphaCnpArrived) {
        dcqcn.alpha = (1 - dcqcn.g) * dcqcn.alpha + dcqcn.g;
    } else {
        dcqcn.alpha = (1 - dcqcn.g) * dcqcn.alpha;
    }
    dcqcn.alphaCnpArrived = false;
    scheduleTimeout(timers.alphaTimer, dcqcn.alphaUpdateInter);
}

void DcqcnRoceConnection::rateInc() {
    //
    if (dcqcn.rateIncStage < 5) { // fast recovery stage
        dcqcn.curRate = (dcqcn.curRate + dcqcn.targetRate) / 2;
    } else if (dcqcn.rateIncStage == 5) { // additive increase
        dcqcn.targetRate += dcqcn.additiveRate;
        if (dcqcn.targetRate > dcqcn.maxRate) // cannot exceed the line rate
            dcqcn.targetRate = dcqcn.maxRate;
        dcqcn.curRate = (dcqcn.targetRate + dcqcn.curRate) / 2;
    } else { // hyper increase
        dcqcn.targetRate += dcqcn.hyperIncRate;
        if (dcqcn.targetRate > dcqcn.maxRate)
            dcqcn.targetRate = dcqcn.maxRate;
        dcqcn.curRate = (dcqcn.targetRate + dcqcn.curRate) / 2;
    }
    dcqcn.rateIncStage += 1;
    scheduleTimeout(timers.rateIncTimer, dcqcn.rateIncInter);
}

void DcqcnRoceConnection::rateDecrease() {
    //
    //
    if (dcqcn.decreaseCnpArrived) {
        dcqcn.decreaseCnpArrived = false;
        dcqcn.targetRate = dcqcn.curRate;
        double rate = dcqcn.curRate * (1 - dcqcn.alpha / 2);
        dcqcn.curRate = std::max(dcqcn.minRate, rate);
        dcqcn.rateIncStage = 0;
        // std::cout << "time: " << omnetpp::simTime() << " current rate: " << dcqcn.curRate
        //           << " dst: " << qp.dstAddr << " src: " << qp.srcAddr << endl;
    }
    scheduleTimeout(timers.rateDecTimer, dcqcn.rateDecInter);
}


void DcqcnRoceConnection::scheduleTimeout(cMessage *msg, omnetpp::simtime_t timeout) {
    scheduleAt(omnetpp::simTime() + timeout, msg);
}

void DcqcnRoceConnection::handleMessage(cMessage *msg) {
    if (msg->isSelfMessage()) {
        if (msg == timers.alphaTimer) {
            updateAlpha();
        } else if (msg == timers.rateIncTimer) {
            rateInc();
        } else if (msg == timers.rateDecTimer) {
            rateDecrease();
        }
    }
}

void DcqcnRoceConnection::handleAck(uint32_t seqNum) {
    senderMgr.handleAck(seqNum);
}

/*
 * Perform Go-back-N logic
 */
void DcqcnRoceConnection::handleNack(uint32_t seqNum) {
    senderMgr.handleNack(seqNum);
}

/*
 * wrap packet
 */

inet::Packet* DcqcnRoceConnection::wrapPacket(
    uint32_t dataSize, CLUSTER_MSGS kind, uint32_t seqNum, 
    bool isLastSegment, const Flow &flow_in) 
{
    if (kind == CLUSTER_MSGS::SOCK_TO_SOCK_DATA_MSG) {
        const auto& roceMsg = makeShared<RoceMsg>();
        roceMsg->setKind(CLUSTER_MSGS::SOCK_TO_SOCK_DATA_MSG);
        roceMsg->setSequenceNum(seqNum);
        roceMsg->setThe_flow(flow_in);
        roceMsg->setIsLastSegment(isLastSegment);
        roceMsg->setChunkLength(inet::B(dataSize));
        std::string flowIndex =string("RoceData-")+std::to_string((int)flow_in.flowId);
        Packet* packet = new Packet(flowIndex.c_str());
        packet->insertAtBack(roceMsg);
        ASSERT(packet != nullptr);
        return packet;
    } else if (kind == CLUSTER_MSGS::ACK_MSG) {
        auto roceMsg = makeShared<RoceMsg>();
        roceMsg->setKind(CLUSTER_MSGS::ACK_MSG);
        roceMsg->setSequenceNum(seqNum);
        roceMsg->setThe_flow(flow_in);
        roceMsg->setChunkLength(inet::B(dataSize));
        roceMsg->setIsLastSegment(isLastSegment);
        Packet* packet = new Packet("RoceACK");
        packet->insertAtBack(roceMsg);
        return packet;
    } else if (kind == CLUSTER_MSGS::NACK_MSG) {
        auto roceMsg = makeShared<RoceMsg>();
        roceMsg->setKind(CLUSTER_MSGS::NACK_MSG);
        roceMsg->setSequenceNum(seqNum);
        roceMsg->setThe_flow(flow_in);
        roceMsg->setChunkLength(inet::B(dataSize));
        roceMsg->setIsLastSegment(isLastSegment);
        Packet* packet = new Packet("RoceNACK");
        packet->insertAtBack(roceMsg);
        return packet;
    } else if (kind == CLUSTER_MSGS::CNP_MSG) {
        auto roceMsg = makeShared<RoceMsg>();
        roceMsg->setKind(CLUSTER_MSGS::CNP_MSG);
        roceMsg->setSequenceNum(-1);
        roceMsg->setThe_flow(flow_in);
        roceMsg->setChunkLength(inet::B(dataSize));
        roceMsg->setIsLastSegment(isLastSegment);
        Packet* packet = new Packet("RoceCNP");
        packet->insertAtBack(roceMsg);
        return packet;
    } else {
        error("DCQCN: Unrecognized message type");
    }
    return nullptr;
}  // namespace ai_sim

}
